#include "path/path.h"

#include <nav_msgs/OccupancyGrid.h>
#include <visualization_msgs/Marker.h>
#include <stdlib.h>

CellNode makeNode(int x, int y, int cost) {
	CellNode node;
	node.x = x;
	node.y = y;
	node.cost = cost;
	return node;
}

Path::Path() :
	robotDiameter_(0.23f),
	targetTag_(NULL)
{
	ros::NodeHandle nh;
	map_publisher_ = nh.advertise<nav_msgs::OccupancyGrid>("/inflated_map", 1);
	path_publisher_ = nh.advertise<visualization_msgs::Marker>("/slam/path", 1);


	// ROS_ERROR("Testing");

	// CellNode n1 = makeNode(1, 1, 0);
	// CellNode n2 = makeNode(1, 1, 1);
	// CellNode n3 = makeNode(1, 1, 2);

	// openCellsQueue_.push(n2);
	// openCellsQueue_.push(n1);
	// openCellsQueue_.push(n3);

	// CellNode n = openCellsQueue_.top(); openCellsQueue_.pop();
	// ROS_ERROR("Popped %d %d %d", n.x, n.y, n.cost);
	// n = openCellsQueue_.top(); openCellsQueue_.pop();
	// ROS_ERROR("Popped %d %d %d", n.x, n.y, n.cost);
	// n = openCellsQueue_.top(); openCellsQueue_.pop();
	// ROS_ERROR("Popped %d %d %d", n.x, n.y, n.cost);

	// closedCells_.insert(n1);
	// if (closedCells_.find(n2) == closedCells_.end()) {
	// 	ROS_ERROR("Test failed");
	// } else {
	// 	ROS_ERROR("Test OK");
	// }
}

Path::~Path() {

}

void Path::setCurrentPosition(float x, float y, float theta) {
	x_ = x;
	y_ = y;
	theta_ = theta;

	// int deleteUpToIndex = -1;
	// std::vector<Waypoint>::iterator it = waypoints_.begin();
	// for (; it < waypoints_.end(); ++it) {
	// 	float wx = (*it).x - x_;
	// 	float wy = (*it).y - y_;
	// 	float dist = sqrtf(wx * wx + wy * wy);

	// 	if (dist < 0.07f) {
	// 		deleteUpToIndex = it - waypoints_.begin();
	// 	}
	// }

	// if (deleteUpToIndex > 0) {
	// 	waypoints_.erase(waypoints_.begin(), waypoints_.begin() + deleteUpToIndex);
	// }
}

void Path::setTarget(TargetTag *tag) {
	targetTag_ = tag;
	if (targetTag_ != NULL) {
		updatePath();
	}
}

// TODO leaks memory if set more than once
void Path::setOccupancyGrid(OccupancyGrid *grid) {
	grid_ = grid;
	width_ = grid_->getWidth();
	height_ = grid_->getHeight();
	resolution_ = grid->getResolution();

	// Pre calculate the inflation region
	inflateRegionSize_ = robotDiameter_ / resolution_ + 0.5f;
	inflateRegion_ = new char*[inflateRegionSize_];
	for (int x = 0; x < inflateRegionSize_; x++) {
		inflateRegion_[x] = new char[inflateRegionSize_];
		for (int y = 0; y < inflateRegionSize_; y++) {
			float centerXDist = x - (inflateRegionSize_ / 2.0f);
			float centerYDist = y - (inflateRegionSize_ / 2.0f);
			float centerDist = sqrtf(centerXDist * centerXDist + centerYDist * centerYDist);
			inflateRegion_[x][y] = centerDist <= (robotDiameter_ / 2.0f);
		}
	}

	cells_ = new char*[width_];
	cellsInflated_ = new char*[width_];
	for (int x = 0; x < width_; x++) {
		cells_[x] = new char[height_];
		cellsInflated_[x] = new char[height_];
	}

	updateRegion(0, 0, width_, height_);
}

void Path::update() {
	DirtyRect dirty = grid_->getDirtyRect();
	grid_->clearDirtyRect();

	updateRegion(dirty.lower_x, dirty.lower_y, dirty.upper_x, dirty.upper_y);
}

void Path::updateRegion(int lowerX, int lowerY, int upperX, int upperY) {
	lowerX = lowerX < 0 ? 0 : lowerX;
	lowerY = lowerY < 0 ? 0 : lowerY;
	upperX = upperX >= width_ ? width_ - 1 : upperX;
	upperY = upperY >= height_ ? height_ - 1 : upperY;

	for (int x = lowerX; x < upperX; x++) {
		for (int y = lowerY; y < upperY; y++) {
			float prob = grid_->probabilityOccupied(x, y);
			cells_[x][y] = prob > 0.70f;
		}
	}

	// Inflate the updated area
	inflateRegion(lowerX , lowerY, upperX, upperY);
}

void Path::inflateRegion(int lowerX, int lowerY, int upperX, int upperY) {
	// Add margin to inflated area
	int halfRegionSize = (inflateRegionSize_ / 2);
	lowerX -= halfRegionSize;
	lowerY -= halfRegionSize;
	upperX += halfRegionSize;
	upperY += halfRegionSize;

	lowerX = lowerX < 0 ? 0 : lowerX;
	lowerY = lowerY < 0 ? 0 : lowerY;
	upperX = upperX >= width_ ? width_ - 1 : upperX;
	upperY = upperY >= height_ ? height_ - 1 : upperY;

	// Clear this region in the inflation cells
	for (int x = lowerX; x < upperX; x++) {
		for (int y = lowerY; y < upperY; y++) {
			cellsInflated_[x][y] = 0;
		}
	}

	lowerX -= halfRegionSize;
	lowerY -= halfRegionSize;
	upperX += halfRegionSize;
	upperY += halfRegionSize;

	lowerX = lowerX < 0 ? 0 : lowerX;
	lowerY = lowerY < 0 ? 0 : lowerY;
	upperX = upperX >= width_ ? width_ - 1 : upperX;
	upperY = upperY >= height_ ? height_ - 1 : upperY;

	// Reinflate from cells_
	for (int x = lowerX; x < upperX; x++) {
		for (int y = lowerY; y < upperY; y++) {
			// Inflate the cell if it is occupied
			if (cells_[x][y]) {
				for (int rx = 0; rx < inflateRegionSize_; rx++) {
					for (int ry = 0; ry < inflateRegionSize_; ry++) {
						int cellx = rx + x - halfRegionSize;
						int celly = ry + y - halfRegionSize;
						if (cellx < 0 || cellx >= width_ || celly < 0 || celly >= height_) continue;
						cellsInflated_[cellx][celly] = 1; //inflateRegion_[rx][ry];
					}
				}
			}
		}
	}
}

void Path::updatePath() {
	waypoints_.clear();
	currentPath_.clear();

	if (targetTag_ == NULL) return;

	// Update path
	int startCellX = x_ / resolution_ + width_ / 2;
	int startCellY = y_ / resolution_ + height_ / 2;
	search(startCellX, startCellY, targetTag_);

	Waypoint waypoint;
	waypoint.x = (*currentPath_.begin()).x * resolution_ - (resolution_ * width_ / 2.0f);
	waypoint.y = (*currentPath_.begin()).y * resolution_ - (resolution_ * height_ / 2.0f);
	waypoints_.push_back(waypoint);

	std::list<CellNode>::iterator it = currentPath_.begin();
	CellNode current = *it;
	it++; // advance one
	CellNode next = *it;
	bool currentDirectionHorizontal = current.y - next.y == 0;
	for (; it != currentPath_.end(); it++) {
		current = next;
		next = *it;
		if (next.x == current.x && next.y == current.y) continue;

		bool newDirectionHorizontal = current.y - next.y == 0;
		if (currentDirectionHorizontal != newDirectionHorizontal) {
			// Direction changed, add waypoint
			waypoint.x = current.x * resolution_ - (resolution_ * width_ / 2.0f);
			waypoint.y = current.y * resolution_ - (resolution_ * height_ / 2.0f);
			waypoints_.push_back(waypoint);
		}
		currentDirectionHorizontal = newDirectionHorizontal;
	}
	// Add last path
	waypoint.x = currentPath_.back().x * resolution_ - (resolution_ * width_ / 2.0f);
	waypoint.y = currentPath_.back().y * resolution_ - (resolution_ * height_ / 2.0f);
	waypoints_.push_back(waypoint);
}

int Path::getNumWaypoints() {
	return waypoints_.size();
}

Waypoint Path::getFirstWaypoint() {
	if (waypoints_.size() > 0) return *waypoints_.begin();
	Waypoint p;
	return p;
}

int travelCost(CellNode from, CellNode to) {
	return abs(from.x - to.x) + abs(from.y - to.y);
}

void Path::search(int fromX, int fromY, TargetTag *tag) {
	closedCells_.clear();
	parentCells_.clear();
	openCells_.clear();
	openCellsQueue_ = std::priority_queue<CellNode, std::vector<CellNode>, CostCellCompare>();
	gScore_.clear();
	fScore_.clear();

	CellNode targetCell = findCellClosestToTag(tag);
	CellNode startCell = makeNode(fromX, fromY, 0.0f);
	targetCell.cost = travelCost(startCell, targetCell);

	openCellsQueue_.push(startCell);
	openCells_.insert(startCell);

	//ROS_ERROR("Starting search from %d %d to %d %d", startCell.x, startCell.y, targetCell.x, targetCell.y);
	gScore_[startCell] = 0;
	fScore_[startCell] = travelCost(startCell, targetCell);

	while (!openCellsQueue_.empty()) {
		CellNode current = openCellsQueue_.top();
		openCellsQueue_.pop();
		openCells_.erase(current);
		CellNode parent = parentCells_[current];
		//ROS_ERROR("Examining cell %d %d", current.x, current.y);

		if (current.x == targetCell.x && current.y == targetCell.y) {
			//ROS_ERROR("Found path!");
			currentPath_.clear();
			// done
			updatePath(current, startCell);

			return;
		}

		closedCells_.insert(current);
		int diffs[4][2] = {{0, 1}, {0, -1}, {1, 0}, {-1, 0}};
		int currentCost = gScore_[current];
		for (int i = 0; i < 4; i++) {
			CellNode neighbor = makeNode(current.x + diffs[i][0], current.y + diffs[i][1], 0);
			if (neighbor.x < 0 || neighbor.x >= width_ || neighbor.y < 0 || neighbor.y >= height_) continue;
			if (closedCells_.find(neighbor) != closedCells_.end()) continue;
			if (cellsInflated_[neighbor.x][neighbor.y]) continue;

			int extraCost = 0;
			if (diffs[i][0] != (parent.x - current.x) && diffs[i][1] != (parent.y - current.y)) {
				extraCost = 5;
			}

			int neighborCost = extraCost + currentCost + travelCost(neighbor, current);
			neighbor.cost = neighborCost;
			//ROS_ERROR("Examining neighbor %d %d %d", neighbor.x, neighbor.y, neighborCost);
			int gScoreNeighbor = 1000000;
			if (gScore_.find(neighbor) != gScore_.end()) {
				gScoreNeighbor = gScore_[neighbor];
			}
			bool neighborInOpenset = openCells_.find(neighbor) != openCells_.end();
			if (!neighborInOpenset || neighborCost <= gScoreNeighbor) {
				parentCells_[neighbor] = current;
				gScore_[neighbor] = neighborCost;
				neighbor.cost = neighborCost + travelCost(neighbor, targetCell);
				fScore_[neighbor] = neighbor.cost;

				if (!neighborInOpenset) {
					//ROS_ERROR("Adding to openset");
					openCells_.insert(neighbor);
					openCellsQueue_.push(neighbor);
				}
			}
		}
	}

	//ROS_ERROR("Search failed");
}

void Path::updatePath(CellNode current, CellNode goal) {
	if (current.x == goal.x && current.y == goal.y) {
		return;
	} else {
		currentPath_.push_front(current);
		updatePath((*(parentCells_.find(current))).second, goal);
	}
}

CellNode Path::findCellClosestToTag(TargetTag *tag) {
	//ROS_ERROR("Looking for cell closest to tag %f %f", tag->getX(), tag->getY());
	int tagX = tag->getX() / resolution_ + width_ / 2;
	int tagY = tag->getY() / resolution_ + height_ / 2;
	int size = inflateRegionSize_ + 4;

	float tx = cosf(tag->getAngle());
	float ty = sinf(tag->getAngle());
	float tlen = sqrtf(tx * tx + ty * ty) / 3.0f;

	float closestDist = 1000000.0f;
	int closestX = 0;
	int closestY = 0;
	for (int x = tagX - size / 2; x < tagX + size / 2; x++) {
		for (int y = tagY - size / 2; y < tagY + size / 2; y++) {
			if (cellsInflated_[x][y]) continue;

			float cx = x - tagX;
			float cy = y - tagY;
			float clen = sqrtf(cx * cx + cy * cy);
			if (clen < closestDist) {
				closestX = x;
				closestY = y;
				closestDist = clen;
			}
		}
	}

	//ROS_ERROR("Found cell %d %d", closestX, closestY);

	return makeNode(closestX, closestY, 0);
}

void Path::renderInflatedMap() {
	std::vector<int8_t> map_data(height_ * width_);
	int count = 0;
	for (int y = 0; y < height_; y++) {
		for (int x = 0; x < width_; x++) {
			int8_t val = cellsInflated_[x][y] ? 100 : 0;
			map_data[count++] = val;
		}
	}

	nav_msgs::OccupancyGrid occupancyGrid;

	occupancyGrid.header.frame_id = "/floor";
	occupancyGrid.header.stamp = ros::Time::now();

	occupancyGrid.info.map_load_time = ros::Time::now();
	occupancyGrid.info.resolution = resolution_;
	occupancyGrid.info.width = width_;
	occupancyGrid.info.height = height_;

	occupancyGrid.info.origin.position.x = -(width_ * resolution_) / 2.0f;
	occupancyGrid.info.origin.position.y = -(height_ * resolution_) / 2.0f;
	occupancyGrid.info.origin.position.z = 0.001f;

	occupancyGrid.info.origin.orientation.x = 0.0f;
	occupancyGrid.info.origin.orientation.y = 0.0f;
	occupancyGrid.info.origin.orientation.z = 0.0f;
	occupancyGrid.info.origin.orientation.w = 1.0f;

	occupancyGrid.data = map_data;

	map_publisher_.publish(occupancyGrid);
}

void Path::renderPath()
{
	visualization_msgs::Marker marker;

	marker.header.frame_id = "/floor";
	marker.header.stamp = ros::Time();

	marker.ns = "path";
	marker.id = 0;
	marker.lifetime = ros::Duration(1.5f);

	marker.type = visualization_msgs::Marker::LINE_STRIP;
	marker.action = visualization_msgs::Marker::ADD;

	marker.pose.position.x = 0.0f;
	marker.pose.position.y = 0.0f;
	marker.pose.position.z = 0.01f;

	marker.pose.orientation.x = 0.0f;
	marker.pose.orientation.y = 0.0f;
	marker.pose.orientation.z = 0.0f;
	marker.pose.orientation.w = 1.0f;

	std::vector<geometry_msgs::Point> points(waypoints_.size());

	std::vector<Waypoint>::iterator it;
	int count = 0;
	for (it = waypoints_.begin(); it < waypoints_.end(); ++it, count++) {
		Waypoint wp = *it;
		points[count].x = wp.x; //node.x * resolution_ - (width_ * resolution_) / 2.0f;
		points[count].y = wp.y; //node.y * resolution_ - (height_ * resolution_) / 2.0f;
		points[count].z = 0.0f;
	}

	marker.points = points;

	marker.scale.x = 0.01f;
	marker.scale.y = 0.1f;
	marker.scale.z = 0.1f;

	marker.color.r = 0.2f;
	marker.color.g = 0.2f;
	marker.color.b = 0.8f;
	marker.color.a = 1.0f;

	// Publish the ray
	path_publisher_.publish(marker);
}
